import numpy as np
import matplotlib.pyplot as plt
from BernoulliBandit import Solver, bandit_10_arm, plot_results

class UCB(Solver):
    def __init__(self, bandit, coef, init_prob=1.0):
        super(UCB, self).__init__(bandit)
        self.total_count = 0
        self.estimates = np.array([init_prob] * self.bandit.K)
        self.coef = coef

    def run_one_step(self):
        self.total_count += 1
        ucb = self.estimates + self.coef * np.sqrt(
            np.log(self.total_count) / (2 * (self.counts + 1)))  # 计算上置信界
        k = np.argmax(ucb)  # 选出上置信界最大的拉杆
        r = self.bandit.step(k)
        self.estimates[k] += 1. / (self.counts[k] + 1) * (r - self.estimates[k])
        return k

def run_single():
    np.random.seed(1)
    ucb_solver = UCB(bandit_10_arm, coef=1)
    ucb_solver.run(5000)
    print('上置信界算法的累积懊悔为：', ucb_solver.regret)
    plot_results([ucb_solver], ["UCB"])

if __name__ == "__main__":
    run_single()
